#!/usr/bin/python3
# coding=utf-8

# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================
import os
import sys
import csv
import time
import logging
import subprocess

import numpy as np

sys.path.append("../..")
from common_scripts.gen_data import MatmulGenData
from common_scripts.compare_data import compare_data
from common_scripts.exec_utils import get_file_work_dir, get_case_list, get_perf_task_duration, clear_file_cache,\
    get_process_case_cmd

IS_BIAS = False
IS_TRANS_A = False
IS_TRANS_B = False
PERF_RUN_TIMES = 5
# float16 in and float32 out
DATA_TYPE_STR = "float16_float32"

logging.basicConfig(level=logging.INFO)


class ProcessParams:
    def __init__(self, case_name, m, n, k, b, is_perf, run_mode):
        self.case_name = case_name
        self.m = m
        self.n = n
        self.k = k
        self.b = b
        self.is_perf = is_perf
        self.run_mode = run_mode


def process_case(file_work_dir, process_params):
    case_name, m, n, k, b, is_perf, run_mode = process_params.case_name, process_params.m, process_params.n,\
        process_params.k, process_params.b, process_params.is_perf, process_params.run_mode
    logging.info("[INFO] start process case [%s]" % (case_name))
    logging.info("[INFO] IS_PERF is set [%d]" % (is_perf))

    loop_num = PERF_RUN_TIMES if is_perf else 1
    task_duration_avg = 0
    for i in range(loop_num):
        clear_file_cache(file_work_dir)
        matmul_gen_data = MatmulGenData(m, n, k, b, IS_TRANS_A, IS_TRANS_B, IS_BIAS, DATA_TYPE_STR)
        if is_perf:
            matmul_gen_data.gen_fake_golden_data(file_work_dir)
        else:
            matmul_gen_data.gen_golden_data(file_work_dir)
        params_str = f"{m} {n} {k}"
        cmd = get_process_case_cmd(kernel_name="ascendc_matmul_ibshareb_bbit", params_str=params_str,\
            is_perf=is_perf, run_mode=run_mode)
        subprocess.run(cmd)
        if is_perf:
            wrong_num = -1
        else:
            logging.info("[INFO] compare data case[%s]" % (case_name))
            wrong_num = compare_data(file_work_dir, DATA_TYPE_STR)
        if is_perf:
            cur_duration = float(get_perf_task_duration("./prof_out"))
            logging.info("[INFO] case_name: [%s], loop_num: [%s], cur_duration: [%s]\n" % (case_name, i, cur_duration))
            task_duration_avg += cur_duration

    res_data = []
    res_data.append(case_name)
    res_data.append(wrong_num)
    res_data.append(b * m * n)
    if wrong_num == -1:
        res_data.append("None")
    elif wrong_num / (b * m * n) > 0.001:
        res_data.append("Fail")
    else:
        res_data.append("Success")
    if is_perf:
        task_duration_avg = task_duration_avg / loop_num
        res_data.append(task_duration_avg)
    return res_data


def main():
    args_len = len(sys.argv) - 1
    if args_len != 2:
        logging.info("[ERROR] exec_test input params error!")
        return -1

    file_work_dir = get_file_work_dir()
    if not os.path.exists(file_work_dir):
        logging.info("[ERROR] file path %s not exist!" % (file_work_dir))
        return -1

    is_perf = False
    if sys.argv[2] == "perf":
        is_perf = True

    case_list = get_case_list()
    res_list = [["case_name", "wrong_num", "total_num", "result", "task_duration"]]
    run_mode = sys.argv[1]
    for is_process, case_name, m, n, k, b in case_list:
        if is_process == 1:
            process_params = ProcessParams(case_name, m, n, k, b, is_perf, run_mode)
            res_data = process_case(file_work_dir, process_params)
            res_list.append(res_data)

    timestamp = time.time()
    result_file_name = "result_" + str(timestamp) + ".csv"
    with open(os.path.join(file_work_dir, "output", result_file_name), 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerows(res_list)

    logging.info("---------------RESULT---------------")
    for res in res_list:
        logging.info(res)
    return 0


if __name__ == "__main__":
    main()
